# api/main.py
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import logging
from pathlib import Path
import asyncio
import base64

from .models import FaceSwapRequest, FaceSwapResponse, TaskQueryResponse, HealthResponse, TaskStatus
from .task_manager import TaskManager
from .faceswap_handler import faceswap_handler
from .cos_config import FaceFusionCOSConfig
from .cos_utils import generate_file_key, upload_to_cos

# 初始化FastAPI应用
app = FastAPI(
    title="FaceFusion API (COS Mode)",
    version="2.1.0",
    description="AI换脸API服务 - COS模式（返回CDN URL）",
    docs_url="/docs",
    redoc_url="/redoc"
)

# CORS配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 全局变量
task_manager: TaskManager = None
logger: logging.Logger = None
cos_client = None
cos_config_info = None

def setup_logging() -> logging.Logger:
    """配置日志系统"""
    log_dir = Path("/root/FaceFusion/logs")
    log_dir.mkdir(parents=True, exist_ok=True)
    log_file = log_dir / "api.log"
    
    logger = logging.getLogger("facefusion_api")
    logger.setLevel(logging.INFO)
    
    if not logger.handlers:
        file_handler = logging.FileHandler(log_file, encoding='utf-8')
        file_handler.setLevel(logging.INFO)
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        formatter = logging.Formatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
        file_handler.setFormatter(formatter)
        console_handler.setFormatter(formatter)
        logger.addHandler(file_handler)
        logger.addHandler(console_handler)
    
    return logger

@app.on_event("startup")
async def startup_event():
    """应用启动事件"""
    global task_manager, logger, cos_client, cos_config_info
    
    logger = setup_logging()
    logger.info("FaceFusion API (COS模式) 启动中...")
    logger.info("固定模式: extreme_beauty (blend=100)")
    
    # 初始化COS
    try:
        cos_config = FaceFusionCOSConfig()
        cos_client = cos_config.get_cos_client()
        cos_config_info = cos_config.get_config_info()
        logger.info(f"COS初始化成功: Bucket={cos_config_info['bucket']}, Region={cos_config_info['region']}")
    except Exception as e:
        logger.error(f"COS初始化失败: {e}")
        logger.warning("将降级为本地模式（返回base64）")
    
    task_manager = TaskManager()
    asyncio.create_task(task_manager.cleanup_old_tasks())
    
    logger.info("FaceFusion API 启动完成")

@app.on_event("shutdown")
async def shutdown_event():
    """应用关闭事件"""
    global logger
    logger.info("FaceFusion API 关闭中...")
    logger.info("FaceFusion API 关闭完成")

async def process_faceswap_task(
    task_id: str,
    source_image_base64: str,
    target_image_base64: str,
    output_quality: int = 100
) -> None:
    """后台处理换脸任务（COS模式）"""
    global task_manager, logger, cos_client, cos_config_info
    
    try:
        logger.info(f"开始处理换脸任务: {task_id} (extreme_beauty模式)")
        
        await task_manager.update_task_status(task_id, TaskStatus.PROCESSING)
        
        result = await faceswap_handler.process_faceswap(
            task_id=task_id,
            source_image_base64=source_image_base64,
            target_image_base64=target_image_base64,
            output_quality=output_quality
        )
        
        if result.success:
            # 上传到COS
            if cos_client and cos_config_info:
                file_key = generate_file_key(task_id, "output.jpg")
                cdn_url = upload_to_cos(
                    cos_client,
                    cos_config_info,
                    result.output_path,
                    file_key
                )
                
                if cdn_url:
                    # 成功上传到COS，保存CDN URL
                    await task_manager.complete_task(task_id, cdn_url)
                    logger.info(f"换脸任务完成(COS): {task_id}, CDN: {cdn_url}")
                else:
                    # COS上传失败，降级为本地路径（后续将转为base64）
                    await task_manager.complete_task(task_id, result.output_path)
                    logger.warning(f"COS上传失败，使用本地模式: {task_id}")
            else:
                # COS未初始化，使用本地模式
                await task_manager.complete_task(task_id, result.output_path)
                logger.warning(f"COS未初始化，使用本地模式: {task_id}")
        else:
            await task_manager.fail_task(task_id, result.error_message)
            logger.error(f"换脸任务失败: {task_id}, 错误: {result.error_message}")
            
    except Exception as e:
        error_msg = f"处理换脸任务时发生未预期错误: {str(e)}"
        await task_manager.fail_task(task_id, error_msg)
        logger.error(f"换脸任务异常: {task_id}, 异常: {str(e)}")

@app.post("/api/faceswap", response_model=FaceSwapResponse, status_code=202, tags=["换脸"])
async def create_faceswap_task(request: FaceSwapRequest) -> FaceSwapResponse:
    """
    创建换脸任务（固定extreme_beauty模式）
    
    - **source_image**: 源图片base64编码
    - **target_image**: 目标图片base64编码  
    - **output_quality**: 输出质量 (1-100，默认100)
    
    返回CDN URL而不是base64
    """
    global task_manager, logger
    
    try:
        logger.info(f"收到换脸任务请求 (extreme_beauty模式, quality={request.output_quality})")
        
        task_id = await task_manager.create_task(
            source_path="temp_source.jpg",
            target_path="temp_target.jpg",
            mode="extreme_beauty",
            output_quality=request.output_quality
        )
        
        asyncio.create_task(
            process_faceswap_task(
                task_id=task_id,
                source_image_base64=request.source_image,
                target_image_base64=request.target_image,
                output_quality=request.output_quality
            )
        )
        
        logger.info(f"换脸任务已创建: {task_id}")
        
        return FaceSwapResponse(
            task_id=task_id,
            message="换脸任务已提交（extreme_beauty模式），请使用task_id查询进度"
        )
        
    except ValueError as e:
        logger.error(f"创建换脸任务参数错误: {str(e)}")
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"创建换脸任务服务器错误: {str(e)}")
        raise HTTPException(status_code=500, detail="服务器内部错误")

@app.get("/api/task/{task_id}", response_model=TaskQueryResponse, tags=["任务"])
async def get_task_status(task_id: str) -> TaskQueryResponse:
    """查询任务状态（返回CDN URL）"""
    global task_manager, logger
    
    try:
        logger.info(f"查询任务状态: {task_id}")
        
        task = await task_manager.get_task(task_id)
        if not task:
            raise HTTPException(status_code=404, detail="任务不存在")
        
        result_url = None
        if task["status"] == TaskStatus.COMPLETED and task["result_path"]:
            result_path = task["result_path"]
            # 判断是CDN URL还是本地路径
            if result_path.startswith("http"):
                # 已经是CDN URL
                result_url = result_path
            else:
                # 本地路径，需要转为base64（降级模式）
                try:
                    with open(result_path, "rb") as f:
                        image_data = base64.b64encode(f.read()).decode('utf-8')
                        result_url = f"data:image/jpeg;base64,{image_data}"
                    logger.warning(f"使用base64降级模式: {task_id}")
                except Exception as e:
                    logger.error(f"读取结果图像失败: {task_id}, 错误: {str(e)}")
        
        return TaskQueryResponse(
            task_id=task["task_id"],
            status=task["status"],
            created_at=task["created_at"],
            updated_at=task["updated_at"],
            result_url=result_url,
            error_message=task["error"]
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"查询任务状态错误: {task_id}, 错误: {str(e)}")
        raise HTTPException(status_code=500, detail="服务器内部错误")

@app.get("/health", response_model=HealthResponse, tags=["Health"])
async def health_check() -> HealthResponse:
    """健康检查接口"""
    return HealthResponse(
        status="healthy",
        gpu_available=True,
        gpu_memory_used="6.2GB",
        gpu_memory_total="15GB",
        active_tasks=0,
        pending_tasks=0
    )

@app.exception_handler(HTTPException)
async def http_exception_handler(request, exc):
    """HTTP异常处理"""
    global logger
    logger.warning(f"HTTP异常: {exc.status_code} - {exc.detail}")
    return {
        "error": True,
        "code": exc.status_code,
        "message": exc.detail
    }
